import torch
import torch.nn as nn


class VGG16(nn.Module):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.features = self.make_layers(in_channels)
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),  # 创建一个全连接层
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, num_classes),
        )

    def make_layers(self, in_channels):
        layers = []
        cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']
        for v in cfg:
            if v == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
                in_channels = v
        # for _ in layers:
        #     print(_)
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)  # nn.Flatten()也可以
        x = self.classifier(x)
        # for layer in self.classifier.children():
        #     print(layer)
        return x


if __name__ == "__main__":
    model = VGG16(in_channels=3, num_classes=10)
    print(f'Total parameters: {sum(param.numel() for param in model.parameters())}')
    for name, param in model.named_parameters():
        print(f"层名: {name}, 参数数量: {param.numel()}")
    x = torch.randn((1, 3, 224, 224))
    preds = model(x)
    print('输出shape:', preds.shape)

# Total parameters: 134309962
# 层名: features.0.weight, 参数数量: 1728                         3 -> 64
# 层名: features.0.bias, 参数数量: 64

# 层名: features.1.weight, 参数数量: 64                          BatchNorm2d(64)
# 层名: features.1.bias, 参数数量: 64

# 层名: features.3.weight, 参数数量: 36864                       64 -> 64
# 层名: features.3.bias, 参数数量: 64

# 层名: features.4.weight, 参数数量: 64
# 层名: features.4.bias, 参数数量: 64

# 层名: features.7.weight, 参数数量: 73728                       64 -> 128
# 层名: features.7.bias, 参数数量: 128

# 层名: features.8.weight, 参数数量: 128
# 层名: features.8.bias, 参数数量: 128

# 层名: features.10.weight, 参数数量: 147456                      128 -> 128
# 层名: features.10.bias, 参数数量: 128

# 层名: features.11.weight, 参数数量: 128
# 层名: features.11.bias, 参数数量: 128

# 层名: features.14.weight, 参数数量: 294912                      128 -> 256
# 层名: features.14.bias, 参数数量: 256

# 层名: features.15.weight, 参数数量: 256
# 层名: features.15.bias, 参数数量: 256

# 层名: features.17.weight, 参数数量: 589824                      256 -> 256
# 层名: features.17.bias, 参数数量: 256

# 层名: features.18.weight, 参数数量: 256
# 层名: features.18.bias, 参数数量: 256

# 层名: features.20.weight, 参数数量: 589824                       256 -> 256
# 层名: features.20.bias, 参数数量: 256

# 层名: features.21.weight, 参数数量: 256
# 层名: features.21.bias, 参数数量: 256

# 层名: features.24.weight, 参数数量: 1179648                      256 -> 512
# 层名: features.24.bias, 参数数量: 512

# 层名: features.25.weight, 参数数量: 512
# 层名: features.25.bias, 参数数量: 512

# 层名: features.27.weight, 参数数量: 2359296                       512 -> 512
# 层名: features.27.bias, 参数数量: 512

# 层名: features.28.weight, 参数数量: 512
# 层名: features.28.bias, 参数数量: 512

# 层名: features.30.weight, 参数数量: 2359296                      512 -> 512
# 层名: features.30.bias, 参数数量: 512

# 层名: features.31.weight, 参数数量: 512
# 层名: features.31.bias, 参数数量: 512

# 层名: features.34.weight, 参数数量: 2359296                        512 -> 512
# 层名: features.34.bias, 参数数量: 512

# 层名: features.35.weight, 参数数量: 512
# 层名: features.35.bias, 参数数量: 512

# 层名: features.37.weight, 参数数量: 2359296                         512 -> 512
# 层名: features.37.bias, 参数数量: 512

# 层名: features.38.weight, 参数数量: 512
# 层名: features.38.bias, 参数数量: 512

# 层名: features.40.weight, 参数数量: 2359296                         512 -> 512
# 层名: features.40.bias, 参数数量: 512

# 层名: features.41.weight, 参数数量: 512
# 层名: features.41.bias, 参数数量: 512

# 层名: classifier.0.weight, 参数数量: 102760448                    25088 -> 4096
# 层名: classifier.0.bias, 参数数量: 4096

# 层名: classifier.3.weight, 参数数量: 16777216                     4096 -> 4096
# 层名: classifier.3.bias, 参数数量: 4096

# 层名: classifier.6.weight, 参数数量: 40960                         4096 ->10
# 层名: classifier.6.bias, 参数数量: 10
# 输出shape: torch.Size([1, 10])
